package com.saml.demo.sp.controller;

import com.alibaba.fastjson.JSON;
import com.saml.demo.sp.constants.SPConstants;
import com.saml.demo.sp.utils.IDPMessageDecoder;
import com.saml.demo.sp.utils.OpenSAMLUtils;
import com.saml.demo.sp.utils.SPCredentials;
import com.saml.demo.sp.vo.UserAccountVo;
import lombok.extern.log4j.Log4j;
import net.shibboleth.utilities.java.support.component.ComponentInitializationException;
import org.opensaml.core.xml.XMLObject;
import org.opensaml.core.xml.schema.impl.XSAnyImpl;
import org.opensaml.messaging.context.MessageContext;
import org.opensaml.messaging.handler.MessageHandler;
import org.opensaml.messaging.handler.MessageHandlerException;
import org.opensaml.messaging.handler.impl.BasicMessageHandlerChain;
import org.opensaml.saml.common.binding.security.impl.MessageLifetimeSecurityHandler;
import org.opensaml.saml.common.messaging.context.SAMLMessageInfoContext;
import org.opensaml.saml.saml2.core.*;
import org.opensaml.saml.saml2.encryption.Decrypter;
import org.opensaml.saml.security.impl.SAMLSignatureProfileValidator;
import org.opensaml.xmlsec.encryption.support.DecryptionException;
import org.opensaml.xmlsec.encryption.support.InlineEncryptedKeyResolver;
import org.opensaml.xmlsec.keyinfo.impl.StaticKeyInfoCredentialResolver;
import org.opensaml.xmlsec.signature.support.SignatureException;
import org.opensaml.xmlsec.signature.support.SignatureValidator;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.servlet.ModelAndView;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * Created by Derry Luo on 2019/5/6.
 */
@Controller
@RequestMapping("/sp/spconsumer")
@Log4j
public class SpConsumerController {

    /**
     * 消费IDP登录返回的response消息
     *
     * @param request
     * @param httpServletResponse
     * @param modelAndView
     * @return
     * @throws Exception
     */
    @RequestMapping("loginconsume")
    public String loginConsume(HttpServletRequest request, HttpServletResponse httpServletResponse, ModelAndView modelAndView) throws Exception {

        log.info("【Get IDP Response】" + JSON.toJSONString(request.getParameterMap()));

        // 从request中解码得到Response
        Response response = (Response) IDPMessageDecoder.decode(request);
        log.info("IDP Response received");
        log.info("IDP Response: ");
        OpenSAMLUtils.logSAMLObject(response);

        //验证目的地址和有效期；
        //validateDestinationAndLifetime(response, request);
        //获得加密后的断言；
        EncryptedAssertion encryptedAssertion = response.getEncryptedAssertions().get(0);
        //获得解密后的断言；
        Assertion assertion = decryptAssertion(encryptedAssertion);
        // TODO 目前这里校验会失败
        verifyAssertionSignature(assertion);
        log.info("Decrypted Assertion: ");
        OpenSAMLUtils.logSAMLObject(assertion);

        logAssertionAttributes(assertion);
        logAuthenticationInstant(assertion);
        logAuthenticationMethod(assertion);

        // 获取用户的邮箱
        String userEmail = this.getAssertionAttribute(assertion, SPConstants.USER_ACCOUNT_ATTR_NAME);
        // 用户的nameId，注销登录时需要使用
        UserAccountVo userAccountVo = new UserAccountVo(userEmail, assertion.getSubject().getNameID().getValue());
        log.info("【user login successfully】 userAccountVo : " + JSON.toJSONString(userAccountVo));

        HttpSession session = request.getSession();
        session.setAttribute(SPConstants.USER_ACCOUNT, userAccountVo);

        redirectToGotoURL(request, httpServletResponse);

        return null;
    }

    /**
     * 消费IDP登出返回的消息
     *
     * @param request
     * @param httpServletResponse
     * @param modelAndView
     * @return
     * @throws Exception
     */
    @RequestMapping("logoutconsume")
    public String logoutConsume(HttpServletRequest request, HttpServletResponse httpServletResponse, ModelAndView modelAndView) throws Exception {

        log.info("【Get IDP Logout Response】" + JSON.toJSONString(request.getParameterMap()));
        log.info("【Http Method】metohd = "+request.getMethod());
        // 从request中解码得到Response
        LogoutResponse response = (LogoutResponse) IDPMessageDecoder.httpRedirectDeflateDecoder(request);
        log.info("IDP Logout Response received: ");
        OpenSAMLUtils.logSAMLObject(response);

        //验证目的地址和有效期；
        //validateDestinationAndLifetime(response, request);
        //获得加密后的断言；
        StatusCode statusCode = response.getStatus().getStatusCode();
        log.info("【validate statusCode】statusCodeValue = " + statusCode.getValue());
        if (StatusCode.SUCCESS.equalsIgnoreCase(statusCode.getValue())) {
            request.getSession().removeAttribute(SPConstants.USER_ACCOUNT);
            return "/index";
        }
        log.error("【Fail to logout.】 message : " + response.getStatus().getStatusMessage());
        return "error";
    }

    /**
     * 解密断言
     *
     * @param encryptedAssertion 加密的断言
     */
    private Assertion decryptAssertion(EncryptedAssertion encryptedAssertion) {
        StaticKeyInfoCredentialResolver keyInfoCredentialResolver
                = new StaticKeyInfoCredentialResolver(SPCredentials.getCredential());

        Decrypter decrypter = new Decrypter(null,
                keyInfoCredentialResolver,
                new InlineEncryptedKeyResolver());
        decrypter.setRootInNewDocument(true);

        try {
            return decrypter.decrypt(encryptedAssertion);
        } catch (DecryptionException e) {
            throw new RuntimeException(e);
        }
    }

    private String getAssertionAttribute(Assertion assertion, String attrName) {
        for (Attribute attribute : assertion.getAttributeStatements().get(0).getAttributes()) {
            if (attribute.getName().equalsIgnoreCase(attrName)) {
                return ((XSAnyImpl) attribute.getAttributeValues().get(0)).getTextContent();
            }
        }
        return null;
    }


    /**
     * 验证目的地址和有效期
     *
     * @param response
     * @param request
     */
    private void validateDestinationAndLifetime(Response response, HttpServletRequest request) {

        MessageContext context = new MessageContext<ArtifactResponse>();
        context.setMessage(response);

        SAMLMessageInfoContext messageInfoContext = context.getSubcontext(SAMLMessageInfoContext.class, true);
        messageInfoContext.setMessageIssueInstant(response.getIssueInstant());

        //生命周期验证，要求SAMLMessageInfoContext包含issue time;
        MessageLifetimeSecurityHandler lifetimeSecurityHandler = new MessageLifetimeSecurityHandler();
        lifetimeSecurityHandler.setClockSkew(1000);
        lifetimeSecurityHandler.setMessageLifetime(2000);
        lifetimeSecurityHandler.setRequiredRule(true);

        //验证消息目的地址，要求base message context包含SAML消息，必需的信息可以从中提取出来
        // TODO 因为验证过程中，通过request.getRequestURL()获取的请求地址是IP:PORT，但是SAML RESPONSE中是域名，无法匹配上，所以不验证地址
//        ReceivedEndpointSecurityHandler receivedEndpointSecurityHandler = new ReceivedEndpointSecurityHandler();
//        receivedEndpointSecurityHandler.setHttpServletRequest(request);
        List handlers = new ArrayList<MessageHandler>();
        handlers.add(lifetimeSecurityHandler);
//        handlers.add(receivedEndpointSecurityHandler);

        BasicMessageHandlerChain<ArtifactResponse> handlerChain = new BasicMessageHandlerChain<ArtifactResponse>();
        handlerChain.setHandlers(handlers);

        try {
            handlerChain.initialize();
            handlerChain.doInvoke(context);
        } catch (ComponentInitializationException e) {
            throw new RuntimeException(e);
        } catch (MessageHandlerException e) {
            throw new RuntimeException(e);
        }
    }


    private void verifyAssertionSignature(Assertion assertion) {

        if (!assertion.isSigned()) {
            throw new RuntimeException("The SAML Assertion was not signed");
        }

        try {
            SAMLSignatureProfileValidator profileValidator = new SAMLSignatureProfileValidator();
            profileValidator.validate(assertion.getSignature());

            SignatureValidator.validate(assertion.getSignature(), SPCredentials.getCredential());

            log.info("SAML Assertion signature verified");
        } catch (SignatureException e) {
            e.printStackTrace();
        }
    }

//    private Credential getCredential() {
//        try {
//            InputStream samlCertificate = SpConsumerController.class.getResourceAsStream("/SamlTest.cer");
//            CertificateFactory certificateFactory = CertificateFactory.getInstance("X.509");
//            //
//            @SuppressWarnings("deprecation")
//            java.security.cert.X509Certificate certificate = (java.security.cert.X509Certificate) certificateFactory.generateCertificate(samlCertificate);
//
//
//            X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec((certificate).getPublicKey().getEncoded());
//            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
//            PublicKey key = keyFactory.generatePublic(publicKeySpec);
//
//            BasicX509Credential credential = new BasicX509Credential(certificate);
////            credential.setPublicKey(key);
//            Object obj = credential.getPublicKey();
//            if (obj instanceof RSAPublicKey) {
//                BigInteger modulus = ((RSAPublicKey) obj).getModulus();
//                BigInteger exponent = ((RSAPublicKey) obj).getPublicExponent();
//                System.out.println("modulus");
//                System.out.println(org.apache.commons.codec.binary.Base64.encodeBase64String(modulus.toByteArray()));
//                System.out.println("public exponent:");
//                System.out.println(org.apache.commons.codec.binary.Base64.encodeBase64String(exponent.toByteArray()));
//
//            }
//            //     System.out.println ("public key is: //n//r"+ credential.getPublicKey());
//            return credential;
//        } catch (Exception e) {
//            log.error(e);
//            throw new RuntimeException(e);
//        }
//    }

    private void redirectToGotoURL(HttpServletRequest req, HttpServletResponse resp) {
        String gotoURL = (String) req.getSession().getAttribute(SPConstants.GOTO_URL_SESSION_ATTRIBUTE);
        log.info("Redirecting to requested URL: " + gotoURL);
        try {
            resp.sendRedirect(gotoURL);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private void logAuthenticationMethod(Assertion assertion) {
        log.info("Authentication method: " + assertion.getAuthnStatements().get(0)
                .getAuthnContext().getAuthnContextClassRef().getAuthnContextClassRef());
    }

    private void logAuthenticationInstant(Assertion assertion) {
        log.info("Authentication instant: " + assertion.getAuthnStatements().get(0).getAuthnInstant());
    }

    private void logAssertionAttributes(Assertion assertion) {
        for (Attribute attribute : assertion.getAttributeStatements().get(0).getAttributes()) {
            log.info("Attribute name: " + attribute.getName());
            for (XMLObject attributeValue : attribute.getAttributeValues()) {
                log.info("Attribute value: " + ((XSAnyImpl) attributeValue).getTextContent());
            }
        }
    }


}
